/*
 * Copyright (c) 2024 Huawei Technologies Co.,Ltd.
 *
 * openGauss is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 * http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 * -------------------------------------------------------------------------
 *
 * htap_helpers.cpp
 * Helper function definitions for HTAP.Holds global configuration for the HTAP module.
 *
 * IDENTIFICATION
 * src/gausskernel/storage/htap/htap_helpers.cpp
 *
 * -------------------------------------------------------------------------
 */

#include "global.h"
#include "funcapi.h"
#include "catalog/pg_operator.h"
#include "nodes/pg_list.h"
#include "nodes/primnodes.h"
#include "utils/date.h"
#include "access/transam.h"
#include "storage/ipc.h"
#include "knl/knl_session.h"
#include "knl/knl_thread.h"
#include "knl/knl_instance.h"
#include "htap_helpers.h"
#include "htap_debug.h"
#include "htap_logger.h"
#include "htap.h"
#include "htap_infra.h"
#include "statistics_manager.h"

#ifndef ENABLE_MOT

uint16_t MOTTimestampToStr(uintptr_t src, char *destBuf, size_t len)
{
    char *tmp = nullptr;
    Timestamp timestamp = DatumGetTimestamp(src);
    tmp = DatumGetCString(DirectFunctionCall1(timestamp_out, timestamp));
    errno_t erc = snprintf_s(destBuf, len, len - 1, tmp);
    pfree_ext(tmp);
    securec_check_ss(erc, "\0", "\0");
    return static_cast<uint16_t>(erc);
}

uint16_t MOTTimestampTzToStr(uintptr_t src, char *destBuf, size_t len)
{
    char *tmp = nullptr;
    TimestampTz timestamp = DatumGetTimestampTz(src);
    tmp = DatumGetCString(DirectFunctionCall1(timestamptz_out, timestamp));
    errno_t erc = snprintf_s(destBuf, len, len - 1, tmp);
    pfree_ext(tmp);
    securec_check_ss(erc, "\0", "\0");
    return static_cast<uint16_t>(erc);
}

uint16_t MOTDateToStr(uintptr_t src, char *destBuf, size_t len)
{
    char *tmp = nullptr;
    DateADT date = DatumGetDateADT(src);
    tmp = DatumGetCString(DirectFunctionCall1(date_out, date));
    errno_t erc = snprintf_s(destBuf, len, len - 1, tmp);
    pfree_ext(tmp);
    securec_check_ss(erc, "\0", "\0");
    return static_cast<uint16_t>(erc);
}

#endif // not defined ENABLE_MOT

namespace htap {
HTAP_DECLARE_LOGGER(HTAPHelpers, HTAP)

// HTAP session
void HTAPDestroySession(MOT::SessionContext *sessionContext)
{
    HTAP_ASSERT(HTAPIsInitialized());
    HTAPInfra *htapInfra = HTAPInfra::GetInstance();
    if (htapInfra != nullptr) {
        HTAP_LOG_TRACE("Destroying session memory context");
        htapInfra->DestroySessionMemoryContext();
    }

    HTAP_LOG_DEBUG("Destroying session context %p, connection_id %u", sessionContext,
        sessionContext->GetConnectionId());
    MOT::GetSessionManager()->DestroySessionContext(sessionContext);
}

// Global map of PG session identification (required for session statistics)
// This approach is safer than saving information in the session context
static pthread_spinlock_t g_sessionDetailsLock;
typedef std::map<MOT::SessionId, pair<::ThreadId, pg_time_t>> SessionDetailsMap;
static SessionDetailsMap sessionDetailsMap;

void HTAPInitSessionDetailsMap()
{
    (void)pthread_spin_init(&g_sessionDetailsLock, 0);
}

void HTAPDestroySessionDetailsMap()
{
    (void)pthread_spin_destroy(&g_sessionDetailsLock);
}

void HTAPRecordSessionDetails()
{
    MOT::SessionId sessionId = u_sess->mot_cxt.session_id;
    if (sessionId != INVALID_SESSION_ID) {
        (void)pthread_spin_lock(&g_sessionDetailsLock);
        (void)sessionDetailsMap.emplace(sessionId, std::make_pair(t_thrd.proc->pid, t_thrd.proc->myStartTime));
        (void)pthread_spin_unlock(&g_sessionDetailsLock);
    }
}

void HTAPClearSessionDetails(MOT::SessionId sessionId)
{
    if (sessionId != INVALID_SESSION_ID) {
        (void)pthread_spin_lock(&g_sessionDetailsLock);
        SessionDetailsMap::iterator itr = sessionDetailsMap.find(sessionId);
        if (itr != sessionDetailsMap.end()) {
            (void)sessionDetailsMap.erase(itr);
        }
        (void)pthread_spin_unlock(&g_sessionDetailsLock);
    }
}

void HTAPClearCurrentSessionDetails()
{
    HTAPClearSessionDetails(u_sess->mot_cxt.session_id);
}

void HTAPGetSessionDetails(MOT::SessionId sessionId, ::ThreadId *gaussSessionId, pg_time_t *sessionStartTime)
{
    // although we have the PGPROC in the user data of the session context, we prefer not to use
    // it due to safety (in some unknown constellation we might hold an invalid pointer)
    // it is much safer to save a copy of the two required fields
    (void)pthread_spin_lock(&g_sessionDetailsLock);
    SessionDetailsMap::iterator itr = sessionDetailsMap.find(sessionId);
    if (itr != sessionDetailsMap.end()) {
        *gaussSessionId = itr->second.first;
        *sessionStartTime = itr->second.second;
    }
    (void)pthread_spin_unlock(&g_sessionDetailsLock);
}

// provide safe session auto-cleanup in case of missing session closure
// This mechanism relies on the fact that when a session ends, eventually its thread is terminated
// ATTENTION: in thread-pooled envelopes this assumption no longer holds true, since the container thread keeps
// running after the session ends, and a session might run each time on a different thread, so we
// disable this feature, instead we use this mechanism to generate thread-ended event into the MOT Engine
static pthread_key_t g_sessionCleanupKey;

static void HTAPSessionCleanup(void *key)
{
    HTAP_ASSERT(!g_instance.attr.attr_common.enable_thread_pool);

    // in order to ensure session-id cleanup for session 0 we use positive values
    MOT::SessionId sessionId = (MOT::SessionId)(((uint64_t)key) - 1);
    if (sessionId != INVALID_SESSION_ID) {
        HTAP_LOG_WARN("Encountered unclosed session %u (missing call to HTAPDestroyTxn()?)", (unsigned)sessionId);
        HTAPClearSessionDetails(sessionId);
        MOT::MOTEngine *engine = MOT::MOTEngine::GetInstance();
        if (engine) {
            MOT::SessionContext *sessionContext = MOT::GetSessionManager()->GetSessionContext(sessionId);
            if (sessionContext != nullptr) {
                HTAPDestroySession(sessionContext);
            }
            // since a call to on_proc_exit(destroyTxn) was probably missing, we should also cleanup thread-locals
            // pay attention that if we got here it means the thread pool is disabled, so we must ensure thread-locals
            // are cleaned up right now. Due to these complexities, onCurrentThreadEnding() was designed to be proof
            // for repeated calls.
            engine->OnCurrentThreadEnding();
        }
    }
}

void HTAPInitSessionCleanup()
{
    (void)pthread_key_create(&g_sessionCleanupKey, HTAPSessionCleanup);
}

void HTAPDestroySessionCleanup()
{
    (void)pthread_key_delete(g_sessionCleanupKey);
}

void HTAPScheduleSessionCleanup()
{
    (void)pthread_setspecific(g_sessionCleanupKey, (const void *)(uint64_t)(u_sess->mot_cxt.session_id + 1));
}

void HTAPCancelSessionCleanup()
{
    (void)pthread_setspecific(g_sessionCleanupKey, nullptr);
}

void HTAPOnSessionClose()
{
    HTAP_LOG_TRACE("Received session close notification (session id: %u, connection id: %u)",
        u_sess->mot_cxt.session_id, u_sess->mot_cxt.connection_id);
    if (u_sess->mot_cxt.session_id != INVALID_SESSION_ID) {
        HTAPClearCurrentSessionDetails();
        if (!HTAPIsInitialized()) {
            HTAP_LOG_ERROR("HTAPOnSessionClose(): HTAP infra is not initialized");
        } else {
            // initialize thread data - this is ok, it will be cleaned up when thread exits
            // avoid throwing errors and ignore them at this phase
            (void)HTAPEnsureSafeThreadAccess();
            MOT::SessionContext *sessionContext = u_sess->mot_cxt.session_context;
            if (sessionContext == nullptr) {
                HTAP_LOG_WARN(
                    "Received session close notification, but no current session is found. Current session id is %u. "
                    "Request ignored.",
                    u_sess->mot_cxt.session_id);
            } else {
                HTAPDestroySession(sessionContext);
            }
        }
    }

    HTAP_ASSERT(u_sess->mot_cxt.session_id == INVALID_SESSION_ID);
    HTAP_ASSERT(u_sess->mot_cxt.session_context == nullptr);
}

/** Notification from thread pool that a pooled thread ended (only when thread pool is ENABLED). */
void HTAPOnThreadShutdown()
{
    // silently ignore if HTAP is not initialized
    HTAPInfra *htapInfra = HTAPInfra::GetInstance();
    if (htapInfra == nullptr) {
        return;
    }

    HTAP_LOG_TRACE("Received thread shutdown notification");
    htapInfra->DestroyThreadMemoryContext();

    MOT::MOTEngine *engine = MOT::MOTEngine::GetInstance();
    if (!engine) {
        HTAP_LOG_ERROR("HTAPOnThreadShutdown(): MOT engine is not initialized");
    } else {
        engine->OnCurrentThreadEnding();
    }
    knl_thread_mot_init(); // reset all thread locals
}

/**
 * on_proc_exit() callback to handle thread-cleanup - regardless of whether thread pool is enabled or not.
 * registration to on_proc_exit() is triggered by first call to EnsureSafeThreadAccess().
 */
void HTAPCleanupThread(int status, Datum ptr)
{
    HTAP_ASSERT(g_instance.attr.attr_common.enable_thread_pool);
    HTAP_LOG_TRACE("Received thread cleanup notification (thread-pool ON, thread id: %u, session id: %u)",
        (unsigned)MOTCurrThreadId, u_sess->mot_cxt.session_id);

    // When FATAL error is thrown, proc_exit doesn't call sess_exit. So we call HTAPOnSessionClose() to make sure
    // session-level resources are destroyed properly.
    HTAPOnSessionClose();

    // when thread pool is used we just cleanup current thread
    // this might be a duplicate because thread pool also calls HTAPOnThreadShutdown() - this is still ok
    // because we guard against repeated calls in MOTEngine::onCurrentThreadEnding()
    HTAPOnThreadShutdown();
}

// in a thread-pooled environment we need to ensure thread-locals are initialized properly
bool HTAPEnsureSafeThreadAccess()
{
    bool threadIdAllocated = false;
    if (MOTCurrThreadId == INVALID_THREAD_ID) {
        HTAP_LOG_DEBUG("Initializing safe thread access for current thread");
        if (MOT::AllocThreadId() == INVALID_THREAD_ID) {
            HTAP_REPORT_ERROR(HTAP_ERROR_RESOURCE_LIMIT, "Thread Init", "Failed to allocate thread identifier");
            return false;
        }
        // register for cleanup only once - not having a current thread id is the safe indicator we never registered
        // proc-exit callback for this thread
        if (g_instance.attr.attr_common.enable_thread_pool) {
            on_proc_exit(HTAPCleanupThread, PointerGetDatum(nullptr));
            HTAP_LOG_DEBUG("Registered current thread for proc-exit callback for thread %p", (void *)pthread_self());
        }
        HTAP_LOG_TRACE("Allocated thread id %u", (unsigned)MOTCurrThreadId);
        threadIdAllocated = true;
    }
    if (MOTCurrentNumaNodeId == MEM_INVALID_NODE) {
        if (!MOT::InitCurrentNumaNodeId()) {
            HTAP_REPORT_ERROR(HTAP_ERROR_RESOURCE_LIMIT, "Thread Init", "Failed to allocate NUMA node identifier");
            return false;
        }
    }
    if (!MOT::InitMasstreeThreadinfo()) {
        HTAP_REPORT_ERROR(HTAP_ERROR_RESOURCE_LIMIT, "Thread Init", "Failed to initialize thread-local masstree info");
        return false;
    }

    // when using thread pool, it is possible that this is a new thread but with an already initialized session
    // in this case we need to allocate statistics slot
    MOT::SessionId sessionId = MOT_GET_CURRENT_SESSION_ID();
    if (g_instance.attr.attr_common.enable_thread_pool && (sessionId != INVALID_SESSION_ID) && threadIdAllocated) {
        if (!MOT::StatisticsManager::GetInstance().ReserveThreadSlot()) {
            HTAP_REPORT_ERROR(HTAP_ERROR_RESOURCE_LIMIT, "Thread Init",
                "Failed to reserve statistics thread slot for new thread %u with existing session %u", MOTCurrThreadId,
                sessionId);
            return false;
        }
    }
    return true;
}

void HTAPDestroyTxn(int status, Datum ptr)
{
    HTAP_ASSERT(!g_instance.attr.attr_common.enable_thread_pool);
    HTAP_LOG_DEBUG("Received thread close notification");

    // cleanup session
    if (!g_instance.attr.attr_common.enable_thread_pool) {
        HTAPCancelSessionCleanup();
    }
    HTAPClearCurrentSessionDetails();
    if (MOT::MOTEngine::GetInstance() == nullptr) {
        HTAP_LOG_ERROR("Cannot destroy transaction: MOT engine is not initialized");
        return;
    }

    MOT::SessionContext *session = (MOT::SessionContext *)DatumGetPointer(ptr);
    if (session != MOT_GET_CURRENT_SESSION_CONTEXT()) {
        if (session == nullptr) {
            HTAP_LOG_WARN("Ignoring request to delete session context: already deleted");
        } else {
            HTAP_LOG_ERROR("Ignoring request to delete session context: incorrect session %u", session->GetSessionId());
        }
    } else if (session != nullptr) {
        MOT_LOG_DEBUG("Destroy SessionContext, connection_id = %u", session->GetConnectionId());
        // initialize thread data, since this call may be accessed from new thread pool worker
        // avoid throwing errors and ignore them at this phase
        (void)HTAPEnsureSafeThreadAccess();
        MOT::GcManager *gc = MOT_GET_CURRENT_SESSION_CONTEXT()->GetTxnManager()->GetGcSession();
        if (gc != nullptr) {
            gc->GcEndTxn();
        }
        HTAPDestroySession(session);
    }

    // clean up thread
    HTAPOnThreadShutdown();
}

static MOT::TxnManager *HTAPInitTxnManager(MOT::ConnectionId connectionId = INVALID_CONNECTION_ID)
{
    if (!u_sess->mot_cxt.txn_manager) {
        bool attachCleanFunc =
            (MOTCurrThreadId == INVALID_THREAD_ID ? true : !g_instance.attr.attr_common.enable_thread_pool);

        // First time we handle this connection
        if (!HTAPIsInitialized()) {
            HTAP_REPORT_ERROR(HTAP_ERROR_INVALID_STATE, "Init Transaction",
                "Cannot create transaction: HTAP module is not initialized");
            return nullptr;
        }

        // create new session context
        MOT::SessionContext *sessionContext =
            MOT::GetSessionManager()->CreateSessionContext(IS_PGXC_COORDINATOR, 0, nullptr, connectionId);
        if (sessionContext == nullptr) {
            HTAP_REPORT_ERROR(HTAP_ERROR_INTERNAL, "Session Initialization", "Failed to create session context");
            return nullptr;
        }
        if (g_instance.attr.attr_common.enable_thread_pool && u_sess->exec_cxt.under_stream_runtime) {
            u_sess->mot_cxt.on_cleanup = HTAPOnSessionClose;
        }
        HTAP_ASSERT(u_sess->mot_cxt.session_context == sessionContext);
        HTAP_ASSERT(u_sess->mot_cxt.session_id == sessionContext->GetSessionId());
        HTAP_ASSERT(u_sess->mot_cxt.connection_id == sessionContext->GetConnectionId());
        HTAP_LOG_TRACE("Created new session (session id: %u, connection id: %u)", u_sess->mot_cxt.session_id,
            u_sess->mot_cxt.connection_id);

        // record session details for statistics report
        HTAPRecordSessionDetails();

        if (attachCleanFunc) {
            // schedule session cleanup when thread pool is not used
            if (!g_instance.attr.attr_common.enable_thread_pool) {
                on_proc_exit(HTAPDestroyTxn, PointerGetDatum(sessionContext));
                HTAPScheduleSessionCleanup();
            } else {
                on_proc_exit(HTAPCleanupThread, PointerGetDatum(nullptr));
                HTAP_LOG_DEBUG("Registered current thread for proc-exit callback for thread %p",
                    (void *)pthread_self());
            }
        }

        u_sess->mot_cxt.txn_manager = sessionContext->GetTxnManager();
        u_sess->mot_cxt.txn_manager->SetIsHtapTxn(true);
        u_sess->mot_cxt.txn_manager->GetGcSession()->SetGcType(MOT::GcManager::GC_TYPE::GC_HTAP_DELTA);
        HTAP_LOG_DEBUG("Initialized transaction manager for thread %" PRIu16, MOTCurrThreadId);
    }

    return u_sess->mot_cxt.txn_manager;
}

/**
 * Creates or retrieves MOT transaction manager.
 * @param txnId transaction id.
 * @return The transaction or null if failed.
 */
MOT::TxnManager *HTAPGetSafeTxn(::TransactionId txnId /* = 0 */)
{
    if (!u_sess->mot_cxt.txn_manager) {
        // result of InitTxnManager() is put in u_sess->mot_cxt.txn_manager, so no need to check return value
        (void)HTAPInitTxnManager();
        if (u_sess->mot_cxt.txn_manager != nullptr) {
            if (txnId != 0) {
                u_sess->mot_cxt.txn_manager->SetTransactionId(txnId);
            }
        } else {
            HTAP_REPORT_ERROR(HTAP_ERROR_INTERNAL, "Init Transaction",
                "Failed to initialize current transaction manager: %s", RcToString(GetRootError()));
        }
    }
    return u_sess->mot_cxt.txn_manager;
}
} // namespace htap
